package org.finesys.chat.core.base.embedding.store.pgvector.handler;

import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import org.finesys.chat.core.base.document.Metadata;
import org.finesys.chat.core.base.embedding.filter.Filter;
import org.finesys.chat.core.base.embedding.store.pgvector.MetadataColumnDefinition;
import org.finesys.chat.core.base.embedding.store.pgvector.MetadataHandler;
import org.finesys.chat.core.base.embedding.store.pgvector.MetadataStorageConfig;
import org.finesys.chat.core.base.embedding.store.pgvector.filter.ColumnFilterMapper;
import org.finesys.chat.core.base.embedding.store.pgvector.filter.PgVectorFilterMapper;
import org.finesys.common.core.util.Utils;
import org.finesys.common.core.util.ValidationUtil;

public class ColumnsMetadataHandler implements MetadataHandler {
    final List<MetadataColumnDefinition> columnsDefinition;
    final List<String> columnsNames;
    final PgVectorFilterMapper filterMapper;
    final List<String> indexes;
    final String indexType;

    public ColumnsMetadataHandler(MetadataStorageConfig config) {
        List<String> columnsDefinitionList = ValidationUtil.ensureNotEmpty(config.columnDefinitions(), "Metadata definition");
        this.columnsDefinition = columnsDefinitionList.stream().map(MetadataColumnDefinition::from).collect(Collectors.toList());
        this.columnsNames = this.columnsDefinition.stream().map(MetadataColumnDefinition::getName).collect(Collectors.toList());
        this.filterMapper = new ColumnFilterMapper();
        this.indexes = Utils.getOrDefault(config.indexes(), Collections.emptyList());
        this.indexType = config.indexType();
    }

    @Override
    public String columnDefinitionsString() {
        return (String) this.columnsDefinition.stream().map(MetadataColumnDefinition::getFullDefinition).collect(Collectors.joining(","));
    }

    @Override
    public void createMetadataIndexes(Statement statement, String table) {
        String indexTypeSql = this.indexType != null ? "USING " + this.indexType : "";
        this.indexes.stream().map(String::trim).forEach(index -> {
            String indexSql = String.format("create index if not exists %s_%s on %s %s ( %s )", table, index, table, indexTypeSql, index);
            try {
                statement.executeUpdate(indexSql);
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
        });
    }

    @Override
    public List<String> columnsNames() {
        return this.columnsNames;
    }

    @Override
    public String whereClause(Filter expression) {
        return this.filterMapper.map(expression);
    }

    @Override
    public Metadata fromResultSet(ResultSet resultSet) {
        Map<String, Object> metadataMap = new HashMap<>();
        Iterator<String> columns = this.columnsNames.iterator();
        while (columns.hasNext()) {
            String columnName = (String) columns.next();
            try {
                Object value = resultSet.getObject(columnName);
                if (value != null) {
                    metadataMap.put(columnName, value);
                }
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
        }
        return new Metadata(metadataMap);
    }

    @Override
    public String insertClause() {
        return (String) this.columnsNames.stream().map((c) -> {
            return String.format("%s = EXCLUDED.%s", c, c);
        }).collect(Collectors.joining(","));
    }

    @Override
    public void setMetadata(PreparedStatement preparedStatement, Integer index, Metadata metadata) {
        int i = 0;
        Iterator<String> columns = this.columnsNames.iterator();
        while (columns.hasNext()) {
            String columnName = (String) columns.next();
            Object value = metadata.get(columnName);
            try {
                preparedStatement.setObject(index + i, value != null ? value : "null");
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
            i++;
        }
    }
}
